'''
Author: Logan.Li
Gitee: https://gitee.com/attacker
email: admin@attacker.club
Date: 2025-03-13 00:44:00
LastEditTime: 2025-03-30 10:04:31
Description: 
'''
import boto3
import time

# 设置区域名称
REGION = 'ap-northeast-1'
# EKS 集群名称
CLUSTER_NAME = 'pro-eks-serverless-cluster'

# 初始化 Boto3 客户端
ec2_client = boto3.client('ec2', region_name=REGION)
efs_client = boto3.client('efs', region_name=REGION)
eks_client = boto3.client('eks', region_name=REGION)

# 获取 EKS 集群的 VPC ID
def get_vpc_id(cluster_name):
    try:
        response = eks_client.describe_cluster(name=cluster_name)
        vpc_id = response['cluster']['resourcesVpcConfig']['vpcId']
        print(f"获取 VPC ID 成功: {vpc_id}")
        return vpc_id
    except Exception as e:
        print(f"❌ 获取 VPC ID 失败: {e}")
        raise

# 获取 VPC 的 CIDR 块
def get_vpc_cidr(vpc_id):
    try:
        response = ec2_client.describe_vpcs(VpcIds=[vpc_id])
        cidr_block = response['Vpcs'][0]['CidrBlock']
        print(f"获取 VPC CIDR 成功: {cidr_block}")
        return cidr_block
    except Exception as e:
        print(f"❌ 获取 VPC CIDR 失败: {e}")
        raise

# 获取 VPC 下的子网 ID
def get_subnets(vpc_id):
    try:
        response = ec2_client.describe_subnets(Filters=[{'Name': 'vpc-id', 'Values': [vpc_id]}])
        subnets = [subnet['SubnetId'] for subnet in response['Subnets']]
        print(f"获取子网成功: {subnets}")
        return subnets
    except Exception as e:
        print(f"❌ 获取子网失败: {e}")
        raise

# 创建 EFS 文件系统
def create_efs():
    try:
        response = efs_client.create_file_system(
            PerformanceMode='generalPurpose',
            Encrypted=False,  # 禁用加密
            Tags=[
                {'Key': 'Name', 'Value': f'{CLUSTER_NAME}-efs'},
                {'Key': 'Environment', 'Value': 'Production'}
            ]
        )
        file_system_id = response['FileSystemId']
        print(f"✅ EFS 文件系统创建中，ID: {file_system_id}")
        
        # 等待文件系统变为可用状态
        while True:
            response = efs_client.describe_file_systems(FileSystemId=file_system_id)
            status = response['FileSystems'][0]['LifeCycleState']
            if status == 'available':
                break
            print(f"✅ EFS 文件系统变为可用状态...当前状态: {status}")
            time.sleep(10)
        
        print(f"✅ EFS 文件系统已创建，ID: {file_system_id}")
        return file_system_id
    except Exception as e:
        print(f"✅ 创建 EFS 文件系统失败: {e}")
        raise

# 创建安全组并添加规则
def create_security_group(vpc_id, cidr_range):
    try:
        response = ec2_client.create_security_group(
            GroupName=f'{CLUSTER_NAME}-efs-sg',
            Description='Security group for EFS access',
            VpcId=vpc_id
        )
        security_group_id = response['GroupId']
        ec2_client.authorize_security_group_ingress(
            GroupId=security_group_id,
            IpProtocol='tcp',
            FromPort=2049,
            ToPort=2049,
            CidrIp=cidr_range
        )
        print(f"安全组已创建，ID: {security_group_id}")
        return security_group_id
    except Exception as e:
        print(f"❌ 创建安全组失败: {e}")
        raise

# 创建挂载目标
def create_mount_target(file_system_id, subnet_id, security_group_id):
    try:
        response = efs_client.create_mount_target(
            FileSystemId=file_system_id,
            SubnetId=subnet_id,
            SecurityGroups=[security_group_id]
        )
        mount_target_id = response['MountTargetId']
        print(f"挂载目标已创建，子网 ID: {subnet_id}, 挂载目标 ID: {mount_target_id}")
        # 获取挂载目标的 DNS 名称
        dns_name = f"{file_system_id}.efs.{REGION}.amazonaws.com"
        mount_point = '/opt'  # 根据需要修改挂载点
        mount_command = f"sudo mount -t nfs4 -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport {dns_name}:{mount_point} {mount_point}"
        print(f"请使用以下命令手动挂载 EFS 文件系统：\n{mount_command}")
        return mount_target_id
    except Exception as e:
        print(f"❌ 创建挂载目标失败: {e}")
        raise

# 清理资源（可选）
def cleanup_resources(file_system_id, security_group_id):
    try:
        # 删除 EFS 文件系统
        print(f"删除 EFS 文件系统: {file_system_id}")
        efs_client.delete_file_system(FileSystemId=file_system_id)
        
        # 删除安全组
        print(f"删除安全组: {security_group_id}")
        ec2_client.delete_security_group(GroupId=security_group_id)
    except Exception as e:
        print(f"❌ 清理资源失败: {e}")

def main():
    try:
        # 获取 VPC ID 和 CIDR
        vpc_id = get_vpc_id(CLUSTER_NAME)
        cidr_range = get_vpc_cidr(vpc_id)
        
        # 获取子网
        subnets = get_subnets(vpc_id)
        
        # 创建安全组
        security_group_id = create_security_group(vpc_id, cidr_range)
        
        # 创建 EFS 文件系统
        file_system_id = create_efs()
        
        # 为每个子网创建挂载目标
        for subnet_id in subnets:
            create_mount_target(file_system_id, subnet_id, security_group_id)
        
        # 如果需要测试后清理资源，取消注释以下代码
        # cleanup_resources(file_system_id, security_group_id)
    
    except Exception as e:
        print(f"❌ 脚本执行失败: {e}")

if __name__ == "__main__":
    main()